"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/19 11:25
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    cp_obViewer_capture_result.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    cp_obViewer_capture_result_VER0.1

------------------------------------------------------------------
# @Description:

------------------------------------------------------------------
"""
import os
import cv2
import shutil
from tqdm import tqdm
from natsort import natsorted


def renameDebugImageName(imgDir: str, imgNameList: list, key='indoor_white_wall_alter.png'):
    for idx, img_name in enumerate(tqdm(natsorted(imgNameList), desc='Renaming...')):
        if '_Color_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_color')}"
        elif '_Depth_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_depth')}"
        elif '_IR_LEFT_' in img_name:
            if idx % 2 == 0:
                new_img_name = f"{idx // 2:04d}_{key.replace('_alter', '_speckle-left')}"
            else:
                new_img_name = f"{idx // 2:04d}_{key.replace('_alter', '_ir-left')}"
        elif '_IR_RIGHT_' in img_name:
            if idx % 2 == 0:
                new_img_name = f"{idx // 2:04d}_{key.replace('_alter', '_speckle-right')}"
            else:
                new_img_name = f"{idx // 2:04d}_{key.replace('_alter', '_ir-right')}"
        else:
            new_img_name = img_name

        ori_img_path = os.path.join(imgDir, img_name)
        new_img_path = os.path.join(imgDir, new_img_name)
        shutil.move(ori_img_path, new_img_path)


def renameNormalImageName(imgDir: str, imgNameList: list, key='indoor_white_wall_alter.png'):
    for idx, img_name in enumerate(tqdm(natsorted(imgNameList), desc='Renaming...')):
        if '_Color_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_color')}"
        elif '_Depth_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_depth')}"
        elif '_IR_LEFT_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_speckle-left')}"
        elif '_IR_RIGHT_' in img_name:
            new_img_name = f"{idx:04d}_{key.replace('_alter', '_speckle-right')}"
        else:
            new_img_name = img_name

        ori_img_path = os.path.join(imgDir, img_name)
        new_img_path = os.path.join(imgDir, new_img_name)
        shutil.move(ori_img_path, new_img_path)


def mvCaptureTrainImageToDstDir(srcDir, dstDir, refName='indoor_white_wall', key='_alter', enDynamic=False):
    if not os.path.exists(srcDir):
        print("src directory not found.")

    if not os.path.exists(dstDir):
        os.makedirs(dstDir)

    color_dir_name = 'Color'
    color_dir_path = os.path.join(srcDir, color_dir_name)
    depth_dir_name = 'Depth'
    depth_dir_path = os.path.join(srcDir, depth_dir_name)
    ir_left_dir_name = 'IR_LEFT'
    ir_left_dir_path = os.path.join(srcDir, ir_left_dir_name)
    ir_right_dir_name = 'IR_RIGHT'
    ir_right_dir_path = os.path.join(srcDir, ir_right_dir_name)

    ref_name = refName + key + '.png'

    # get img name
    color_name_list = natsorted(os.listdir(color_dir_path))
    depth_name_list = natsorted(os.listdir(depth_dir_path))
    ir_left_name_list = natsorted(os.listdir(ir_left_dir_path))
    ir_right_name_list = natsorted(os.listdir(ir_right_dir_path))

    # rename
    if enDynamic:
        renameNormalImageName(color_dir_path, color_name_list, key=ref_name)
        renameNormalImageName(depth_dir_path, depth_name_list, key=ref_name)
        renameNormalImageName(ir_left_dir_path, ir_left_name_list, key=ref_name)
        renameNormalImageName(ir_right_dir_path, ir_right_name_list, key=ref_name)
    else:
        renameDebugImageName(color_dir_path, color_name_list, key=ref_name)
        renameDebugImageName(depth_dir_path, depth_name_list, key=ref_name)
        renameDebugImageName(ir_left_dir_path, ir_left_name_list, key=ref_name)
        renameDebugImageName(ir_right_dir_path, ir_right_name_list, key=ref_name)

    # update list
    color_name_list = natsorted(os.listdir(color_dir_path))
    depth_name_list = natsorted(os.listdir(depth_dir_path))
    ir_left_name_list = natsorted(os.listdir(ir_left_dir_path))
    ir_right_name_list = natsorted(os.listdir(ir_right_dir_path))

    color_path_list = [os.path.join(color_dir_path, img_name) for img_name in color_name_list]
    depth_path_list = [os.path.join(depth_dir_path, img_name) for img_name in depth_name_list]
    ir_left_path_list = [os.path.join(ir_left_dir_path, img_name) for img_name in ir_left_name_list]
    ir_right_path_list = [os.path.join(ir_right_dir_path, img_name) for img_name in ir_right_name_list]

    all_ori_img_list = color_path_list + depth_path_list + ir_left_path_list + ir_right_path_list

    for idx, ori_img_path in enumerate(tqdm(natsorted(all_ori_img_list), desc="Processing...")):
        ori_img_name = ori_img_path.split('/')[-1]
        dst_img_path = os.path.join(dstDir, ori_img_name)
        if not os.path.exists(dst_img_path):
            shutil.move(ori_img_path, dst_img_path)


if __name__ == '__main__':
    ori_capture_dir = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/precision/tmp'
    save_dir = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/precision/normal-precision-new'
    mvCaptureTrainImageToDstDir(ori_capture_dir, save_dir, refName='indoor_white_wall', key='_alter', enDynamic=True)
